#!/usr/bin/env python

from numpy.fft import fft
from numpy import real, imag, array, polyfit
from math import atan2, sqrt, pi
from pylab import plot, show
from random import normalvariate

def getPhase(cf):
    cfa = array(cf).astype('d')
    cplx = fft(cfa)
    re = real(cplx)
    #print re
    im = imag(cplx)
    #print im
    phase = []
    for j in range(len(re)):
        phase.append(float(atan2(im[j],re[j])))
    return re,im,phase[0:16]

def getMagnitude(cf):
    cfa = array(cf).astype('d')
    cplx = fft(cfa)
    re = real(cplx)
    im = imag(cplx)
    mag = []
    for j in range(len(re)):
        mag.append(float(sqrt(re[j]**2 + im[j]**2)))
    return mag[0:16]

def dewrap(phase):
    phout = list(phase[:])
    for i in range(1,len(phase)-1):
        diff = phase[i+1]-phase[i]
        #print diff
        if diff>0 and diff>4:
            #print 'positive wrap at '+str(i)+'!'
            wrap = -1
        elif diff<0 and diff<-4:
            #print 'negative wrap at '+str(i)+'!'
            wrap = 1
        else:
            wrap = 0
        for j in range(i+1,len(phase)):
            phout[j] += 2*pi*wrap
    return tuple(phout)

def getPhaseFit(phase):
    [m, b] = polyfit( range(1,len(phase)), phase[1:], 1)
    m = float(m)
    b = float(b)
    fit = range(len(phase))
    for x in fit:
        fit[x] = m*x + b
    return ((m,b), tuple(fit))

def corrReorder(cin):
    cpos = list(cin[16:32])
    cneg = list(cin[0:16])
    cout = cpos
    cout.extend(cneg)
    return cout

def randArray(length, mu=0, sigma=1):
    arr = []
    for i in range(length):
        arr.append(int(normalvariate(mu,sigma)))
    return arr

def crossCorr(a, b, lags=range(-8,8)):
    if len(a)!=len(b):
        print "Vectors must be the same length,"
        print "but vector 'a' has length",len(a)
        print "and vector 'b' has length",len(b)
        return None
    cf = []
    for i in lags:
        if i>0:
            m1 = a[i:]
            m2 = b[:-i]
        elif i==0:
            m1 = a
            m2 = b
        else:
            m1 = a[:i]
            m2 = b[-i:]
        cf.append(sum(m1*m2))
    return array(cf)


